import traceback

from dawate import main_stream
import pymysql
import re
import json_repair
import json
import requests
import sys
def dawate_streaming(content):
    #     ls = []
    #     ls.append({"role": "user", "content": '肖建毅参与的项目有哪些'})
    #     ls.append({"role": "assistant", "content": '''1、2018年，负责组织建设广东电网有限责任公司应用级灾备中心一期建设项目、广东电网公司营销系统维护、广东电网公司资产系统维护、广东电网公司财务系统维护、广东电网公司人资系统维护、广东电网公司GIS维护等项目。
    # 2、2019年，负责组织建设广东电网有限责任公司应用级灾备中心二期建设项目、信息中心众测平台移动应用项目广东电网公司财务系统维护、广东电网公司信息中心1000号服务、广东电网公司营销系统维护、广东电网公司资产系统维护、广东电网公司财务系统维护、广东电网公司人资系统维护等项目。
    # 3、2020年，负责组织建设IOS系统实用化项目及面向电力行业的数据库智能运维模型及应用研究项目、广东电网公司信息中心1000号服务、广东电网公司营销系统维护、广东电网公司资产系统维护、广东电网公司财务系统维护、广东电网公司人资系统维护等项目。
    # 4、2021年，负责组织建设南方电网公司协同办公系统V2.0。'''})
    # for i in range(5):
    #     ls.append({"role": "user", "content": ''.join(["你是谁"]*5000)})
    #     ls.append({"role": "assistant", "content": '我是大瓦特'})

    # for i in ls:
    #     print(i)

    config_params = {
        'relAppId': '17987690',
        'appId': '17987868',
        'appSecret': '5570ea83260dbcad482a7c9175411d5e',
        # 'messages': [{"role": "user", "content": ''.join(["你是谁"]*5000)}]
        'messages': [{"role": "user", "content": content}]
        # 'messages': ls
    }
    a = ''
    response_content = main_stream(config_params)
    try:
        while True:
            a += next(response_content)
            sys.stdout.write(f"\r{[a]}")
            sys.stdout.flush()
    except StopIteration:
            pass
    return a
DIFY_URL = 'http://192.168.73.129:80'
def db_connet():
    conn = pymysql.connect(
        host='127.0.0.1',  # 连接名称，默认127.0.0.1
        user='root',  # 用户名
        passwd='123456',  # 密码
        port=3306,  # 端口，默认为3306
        db='zentao',  # 数据库名称
        charset='utf8',  # 字符编码
    )
    # conn = pymysql.connect(**PYMYSQL_CONF)
    print("----数据库连接成功------")
    # logging.info("----数据库连接成功------ %s", conn)
    return conn
def data_base_connet(sql_statement):
    # 获取游标
    conn=db_connet()
    cursor = conn.cursor()
    # print(cursor)

    # 执行sql语句execute和executemany
    # 定义要执行的SQL语句列表
    try:

        hearders, results = [[], []]

        print(f'--------------------------\n{sql_statement}\n---------------------------')

        cursor.execute(sql_statement)

        #简单用开头是不是select来判断是否为查询
        if re.sub(r'[^\w\u4e00-\u9fff]+', '', sql_statement).lower().startswith('select'):
            hearders = [description[0] for description in cursor.description]
            results = cursor.fetchall()
        else:
            conn.commit()  # 事务的提交

        cursor.close()  # 关闭查询游标
        conn.close()  # 查询完毕，需要关闭连接，释放计算机资源

        return hearders, results
    except pymysql.MySQLError as e:
        print(f"An error occurred: {e}")
        cursor.close()  # 关闭查询游标
        conn.commit()  # 事务的提交
        conn.close()  # 查询完毕，需要关闭连接，释放计算机资源
        # print('sql执行成功')
        return e,None
#获取用户权限
def filter_user_permissions(account):

    #查询用户权限
    hearders = None
    results = None
    sql_statement = f'''
                    SELECT p.* FROM zt_user u left join agent_permisions p on p.role=u.role WHERE account = '{account}';
                    '''
    hearders, results = data_base_connet(sql_statement)

    user_role = results[0][1]

    if not user_role:
        return False, hearders,results

    return True, hearders,results

#json格式化
def json_process(data_answ, pattern=r'```json(.*?)```'):
    matches = re.findall(pattern, data_answ.split('</think>')[-1], re.DOTALL)  # 抽取出内容并判断长度
    # print(matches)

    if '```json' in pattern:
        if matches:
            c_statement = json_repair.loads(matches[-1])
            # print("\n-------step2查到的对应的内容为：\n", c_statement)
        else:
            c_statement = {}
    else:
        if matches:
            c_statement = matches[-1]
            # print("\n-------step2查到的对应的内容为：\n", c_statement)
        else:
            c_statement = ''
    # decoded_object = json_repair.loads(c_statement)
    # return json.dumps(decoded_object,ensure_ascii=False)
    return c_statement
def dify_chat_streaming(query,api_key):
    # API端点和API密钥
    # url = 'http://172.16.32.45:11206/v1/chat-messages' #v1/chat-messages'调用聊天框  workflows/run智能体
    url = 'http://192.168.73.129:80/v1/chat-messages' #v1/chat-messages'调用聊天框  workflows/run智能体
    # api_key = 'app-poRofzsvJiJGhpc1Wn1ecEbg'  # 替换为您的实际API密钥

    # 请求头
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }

    # 请求数据
    data = {
        "inputs": {},
        "query": query,
        "response_mode": "streaming", #blocking streaming
        "conversation_id": "",
        "user": "abc-123",
            }

    # 将数据转换为JSON字符串
    data_json = json.dumps(data)

    # 发送POST请求
    # response = requests.post(url, headers=headers, data=data_json.encode("utf-8"))
    # response.encoding = 'utf-8'

    answer = ''

    # 发送GET请求，并设置stream=True
    with requests.post(url, headers=headers, data=data_json.encode("utf-8"), stream=True) as r:
        # 检查请求是否成功
        r.raise_for_status()
        # print('r',r)
        # if '200'in r:
        #     print(json.loads(r.text))
        # 打印每一行文本
        for line in r.iter_lines():
            # 忽略保持连接的空行
            if line:
                # print(line)
                # 解码每一行，因为返回的数据可能是字节串
                line = line.decode('utf-8')
                if 'data: ' in line:
                    data = json.loads(line.replace('data: ', ''))
                    if data['event'] == "message":
                        print(data['answer'],end='')
                        answer += data['answer']
    return answer
def system_intention(query, old_system,user_account):
    system = ''
    error = ''
    # 查询用户权限
    # #ste0 用户登陆后权限过滤
    print('获取用户权限')
    judge, permission_hearders, permission_results = filter_user_permissions(user_account)
    if not judge:
        # 用户无权限
        error = '\n**用户没有权限**\n'
        return system, error

    # 系统意图识别
    # if not system:
    #     yield '\t\t\n--------------系统意图识别-----------------\n'
    stream_content=dify_chat_streaming(query, api_key="app-xVvAg9zhpVHIeTnXpgzURc3l")
    # stream_content = dawate_streaming(query)

    json_statement = json_process(stream_content)

    # 如果有提示语，则为该系统不支持操作，结束流程
    if json_statement['提示语']:
        if not old_system:
            error = f"\n**{json_statement['提示语']}**\n"
            return system, error
        else:
            return old_system, error

    else:
        if '禅道' in json_statement['业务域意图']:
            system = json_statement['业务域意图']
            return system, error

        else:
            system = json_statement['业务域意图']
            error = f"\n**该系统暂不支持**\n"
            return system, error
#判断操作意图
def operate_intention(query, old_operate):

    operate = ''
    error = ''

    stream_content = dify_chat_streaming(query=query, api_key='app-B1JNuJjuhFC2HHpzW2zqwPHE')

    data_answ=stream_content
    json_statement = json_process(data_answ)
    print('json_statement:',json_statement)
    # 如果为普通问答，则为该系统不支持操作，结束流程
    if json_statement['操作意图'] == '普通问答':
        if not old_operate:
            error = f"\n**系统不支持该操作**\n"
            return operate, error
        else:
            return old_operate, error

    else:
        operate = json_statement['操作意图']
        return operate, error

#安全策略过滤
def Security_Identification():
    Security_prompt = '''
    安全策略过滤：
    要求：
    1。恶意指令，对于危害系统的操作直接禁止，阻断执行,严格禁止的系统的操作行为包括以下内容（不限于以下内容）：
    （1）"删除所有用户" ;
    (2)删除所有数据;
    (3)清空数据库;
    (4)修改所有用户的密码;
    (5)将所有用户设置为管理员;
    (6)向数据库中插入大量无效/垃圾数据;
    (7)随机修改数据库记录;
    (8)删除核心系统文件;
    (9)导出所有用户的敏感信息（如邮箱、电话、住址）” / “列出所有管理员账户
    2.越权操作，动态权限校验，用户只能操作自己拥有的权限，越权直接拒绝；
    3.数据泄露，自动脱敏，查询结果中身份证号显示为****；

    4.注入攻击，SQL语法树分析，检测DROP TABLE语句特征
    '''

    return Security_prompt

#sql转为markdown格式
def sql2md(columns, data, conversation_id=''):
    if len(columns) == 0 or len(data) == 0:
        return f"\n未查询到数据"

    # 开始构建Markdown表格
    md_table = '| ' + ' | '.join(columns) + ' |\n'
    md_table += '| ' + ' | '.join(['---'] * len(columns)) + ' |\n'

    # 添加每一行数据
    for row in data:
        row_fillna = []
        for i in row:
            if i:
                row_fillna.append(str(i))
            else:
                row_fillna.append('/')
        md_table += '| ' + ' | '.join(field for field in row_fillna) + ' |\n'
        # md_table += '| ' + ' | '.join(str(field) for field in row) + ' |\n'

    return f"sql查询结果如下：\n{md_table}\n"

#禅道任务操作
def biz_td(data,user_data):
    user_input = user_data[0]
    intention_user = user_data[1]


    #查询相关表字段，目前只有zt_task，后续可以换成知识库

    #涉及的表
    # tables = ['zt_task']

    sql_statement = '''
                SELECT 
    COLUMN_NAME AS '字段名称',
    COLUMN_COMMENT AS '字段备注',
    COLUMN_COMMENT AS '字段备注'
FROM 
    INFORMATION_SCHEMA.COLUMNS 
WHERE 
    TABLE_NAME = 'zt_task'
    AND TABLE_SCHEMA = 'zentao'; 
        '''
    hearders, results = data_base_connet(sql_statement)

    # 替换项目成员
    sql_1 = '''
                        SELECT account, realname 
                        FROM zt_user 
                        WHERE account IN (
                            SELECT account 
                            FROM zt_team 
                            WHERE root = '1894'
                        );
                    '''
    hearders1, results1 = data_base_connet(sql_1)  # 获取指定项目的团队成员

    # 中英文映射表
    en_zh = []
    mid = {"指派人":[]}
    members = []
    for result in results1:
        mid['指派人'].append({"展示": result[1], '存储': result[0]})
        members.append(result[1])
    en_zh.append(json.dumps(mid, ensure_ascii=False))

    #难度
    difficulty = {'难度':[{"展示": '非常重要', '存储': 'veryImportant'},
                        {"展示": '重要', '存储': 'important'},
                        {"展示": '一般', '存储': 'general'},
                        {"展示": '不重要', '存储': 'notImportant'}]}
    en_zh.append(json.dumps(difficulty, ensure_ascii=False))

    #任务类型
    task_type = {
        '任务类型':[
            {"展示": '设计', '存储': 'design'},
            {"展示": '开发', '存储': 'devel'},
            {"展示": '需求', '存储': 'request'},
            {"展示": '测试', '存储': 'test'},
            {"展示": '实施', '存储': 'implement'},
            {"展示": '研究', '存储': 'study'},
            {"展示": '讨论', '存储': 'discuss'},
            {"展示": '界面', '存储': 'ui'},
            {"展示": '事务', '存储': 'affair'},
            {"展示": '报障', '存储': 'fault'},
            {"展示": '其他', '存储': 'misc'},
        ]}
    en_zh.append(json.dumps(task_type, ensure_ascii=False))
    en_zh = '\n'.join(en_zh)

    #查询意图对应的模板
    intention2temp = {'创建任务': '/zentao/task/create',
                      '修改任务': '/zentao/task/update',
                      '查询任务': '',
                      '删除任务': '',
                      }
    if intention_user in intention2temp.keys():
        full_name = intention2temp[intention_user]

        if full_name:
            # 查询模板
            sql2 = f'''select content  from agent_template where full_name='{full_name}';'''

            hearders2, results2 = data_base_connet(sql2)
            template = results2[0][0]

            #如果是创建和修改任务，需要替换团队成员
            if full_name in ['/zentao/task/create', '/zentao/task/update']:

                template = template.replace('张三/李四', '/'.join(members))
        else:
            template = ''

    else:
        print( '\t\t\n暂不支持该操作\n')
        return


    #安全策略识别
    tem_prompt = f'''
                  安全策略为：{Security_Identification()};
                          当前用户的用户权限为：{data}；
                          当前用户的操作意图为：{intention_user}
                          用户输入为：{user_input};
                          模板为：\n{template}'''
    # print("安全策略权限过滤：", len(tem_prompt))

    # 安全策略识别
    print( '\t\t\n--------------安全策略识别-----------------\n')
    stream_content = dify_chat_streaming(tem_prompt, api_key='app-kq9Bj0KXAPNjYKUyZd0gNoYJ')

    data_answ=stream_content
    st = json_process(data_answ)

    print(f'\n用户的操作合规性判断：{st}')
    print( f"\t\t\n用户的操作合规性判断：{st}\n")

    #如果不合规，直接结束，并返回模板（如果有）
    if st.get('是否合规') != '是':
        print( '\t\t\n用户操作不合规\n')

        #如果有模板，返回模板
        if template:
            print(f'{full_name}模板为:\n{template}')
            print( f'\t\t\n{full_name}模板为:\n{template}\n')

        return

    SQL_prompt = f'''
                    project字段固定为1893，execution固定为1894
                    用户问题为：{user_input}，
                    数据表为：zt_task
                    数据库表字段为：{hearders}{results}
                    部分字段用户输入为中文，但是数据库需要存英文，以下为部分字段的映射关系：
                    {en_zh}
                        '''

    # 生成sql中
    print('\t\t\n--------------生成sql中-----------------\n')
    stream_content = dify_chat_streaming(SQL_prompt, api_key='app-vRtl7zmY68hCnWB1QNWK8Yhn')
    data_answ = ''
    try:
        while True:
            gen_text = next(stream_content)
            gen_text_split = gen_text.split('\t')
            conversation_id = gen_text_split[1]
            answer = '\t'.join(gen_text_split[2:])
            data_answ += answer
            # print(answer, end='')
            print( gen_text)

    except StopIteration as e:
        # conversation_id = e.value
        pass

    # print('生成的sql为：', data_answ)
    # yield f'\t\t\n生成的sql为：\n{data_answ}'

    sql_statement_ = json_process(data_answ, pattern=r'```sql(.*?)```')

    headers, results = data_base_connet(sql_statement_)
    # print('数据库受影响行数:', headers, results)
    # print('完成创建任务')
    # yield '\t\t\n完成创建任务\n'
    #如果有结果，则转为md格式展示，否则为成功
    if results:

        #转为markdown
        print( f"\t\t\n{sql2md(hearders, results)}\n")

    else:
        print( f'\t\t\n完成{full_name}\n')

def ai_qa_ge(query, user_account, operate):
    print('filter_user_permissions')
    judge, permission_hearders, permission_results = filter_user_permissions(user_account)
    print(judge)
    stream_content = biz_td([permission_hearders,permission_results],[query,operate,user_account])

    return stream_content
def run():
    query='禅道系统创建任务'
    old_system=''
    user_account='liuxueping'
    old_operate=''
    # 业务域意图识别
    try:
        old_system,error=system_intention(query, old_system,user_account)
        print(old_system,error)
        if error:
            return
    except:
        print(traceback.format_exc())
        # 大模型错误
        print('大模型错误')
    print("---------------------业务意图",old_system)

    # 操作意图识别
    try:

        old_operate,error=operate_intention(query, old_operate)
        if error:
            return
    except:
        print(traceback.format_exc())
        # 大模型错误
        print('大模型错误')
    print('---------------------操作意图：', old_operate)
    print('----------安全策略识别')
    print(query,  user_account, old_operate)
    query='禅道系统创建任务'
    user_account='liuxueping'
    old_operate="创建任务"
    print('----------安全策略识别')
    ai_qa_ge(query,user_account, old_operate)
run()
# query='禅道系统创建任务'
# user_account='liuxueping'
# old_operate="创建任务"
# print('----------安全策略识别')
# s = ai_qa_ge(query,user_account, old_operate)
# # while True:
#     print(next(s))


